//#include <pthread.h>
#include <mutex>
#include <thread>
//#include <unistd.h>
#include <signal.h>

#ifndef USE_LINUX
#include <winsock2.h>
#endif

#include <mysql.h>
#include <map>
#include "opencv2/opencv.hpp"

#include "cnn_function.hpp"
#include "util.h"

#include "rpc/server.h"

std::vector<std::string> current_user_name;
std::vector<int> current_user_spoofing;
std::vector<cv::Rect> rois_global;

// sudo apt-get install libmysqlclient-dev
std::map<int, float*> all_feature;

bool have_init = false;
std::mutex gpu_lock;
bool recognition_stop = false;

int reload_feature();

void free_dnn_network(){
    uninit_detector();
    uninit_recognition();
    uninit_mtcnn();
}

void init_dnn_network(){
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    reload_feature();
    have_init = true;
    std::string cnn_dir = "/home/cj/face/cnn";
    init_recognition((cnn_dir + "/cfg/cosface_new_34.cfg").c_str(), (cnn_dir + "/model/face_30_glint_refine.cnn").c_str());
    //init_recognition((cnn_dir + "/cfg/cosface_new_100.cfg").c_str(), (cnn_dir + "/model/model.cnn.100").c_str());
    init_detector((cnn_dir + "/cfg/yolov3-small.cfg").c_str(), (cnn_dir + "/model/yolov3-small_final_max_epoch_15.weights").c_str());
    init_mtcnn((cnn_dir + "/cfg/mtcnn_onet.cfg").c_str(), (cnn_dir + "/model/mtcnn_final.weights").c_str());
}

void print_error_message(MYSQL *connection){
    if (mysql_errno(connection)) {
        fprintf(stderr, "mysql error %d: %s\n", mysql_errno(connection), mysql_error(connection));
    }
}

int add_feature(int feature_id, std::string feature_str){
    std::map<int, float*>::iterator it = all_feature.find(feature_id);
    if (it == all_feature.end()){
        const char *float_str = feature_str.c_str();
        float *feature_float = (float *)malloc(512 * sizeof(float));
        for(unsigned int j = 0; j < 512; j++) {
            feature_float[j] = atof(float_str);
            float_str = strchr(float_str, ',') + 1;
        }
        all_feature[feature_id] = feature_float;
    } else {
        printf("add_feature: feature exist\n");
        reload_feature();
    }
    return 0;
}

int reload_feature() {
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    for(std::map<int, float*>::iterator it = all_feature.begin(); it != all_feature.end(); it++){
        free(it->second);
    }
    all_feature.clear();
    MYSQL *connection = mysql_init(NULL);
    if (connection == NULL) {
        printf("mysql_init failed!\n");
        return -1;
    }
    connection = mysql_real_connect(connection, "localhost", "lbtc", "sxfMd4_f12508ccsdfdf", "recognition", 0, NULL, 0);
    if (connection) {
        printf("Connection success! ");
    } else {
        print_error_message(connection);
        mysql_close(connection);
        return 1;
    }
    char query_sql[256] = {0};
    sprintf(query_sql, "select id, feature from person where deleted = 0 order by id");
    int query_result = mysql_query(connection, query_sql);
    if(query_result){
        print_error_message(connection);
        mysql_close(connection);
        return 2;
    } else {
        MYSQL_RES *res_ptr = mysql_store_result(connection);
        if(res_ptr) {
            unsigned long row_num = mysql_num_rows(res_ptr);
            printf("query %lu Rows\n", row_num);
            unsigned int num_fields = mysql_num_fields(res_ptr);
            printf("num_fields: %u\n", num_fields);
            MYSQL_ROW row;
            while((row = mysql_fetch_row(res_ptr)))  {
                //unsigned long *field_lengths = mysql_fetch_lengths(res_ptr);
                int feature_id = -1;
                float *feature_float = (float *)malloc(512 * sizeof(float));
                for(unsigned int i = 0; i < num_fields; i++){
                    if(i == 0){
                        feature_id = atoi(row[i]);
                    } else if (1 == i){
                        char *float_str = (char *)(row[i]);
                        for(unsigned int j = 0; j < 512; j++) {
                            feature_float[j] = atof(float_str);
                            float_str = strchr(float_str, ',') + 1;
                        }
                    } else {
                        printf("reload_feature: database field error\n");
                        exit(-1);
                    }
                }
                all_feature[feature_id] = feature_float;
            }
            if (mysql_errno(connection)) {
                fprintf(stderr, "Retrive error:%s\n", mysql_error(connection));
            }
            mysql_free_result(res_ptr);
        } else {
            print_error_message(connection);
            mysql_close(connection);
            return 3;
        }
    }
    mysql_close(connection);
    printf("all_feature size %lu\n", all_feature.size());
    printf("reload_feature over\n");
    return 0;
}

std::string registration_image(std::vector<char> const& img, int is_recognition) {
    double start = getMsecOfNow();
    cv::Mat image_input = cv::imdecode(img, CV_LOAD_IMAGE_COLOR);
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);

    int code = 0;
    if(!have_init){
        init_dnn_network();
    }
    int detection_bbox[MAX_BBOX_NUM * 4];

    std::vector<cv::Mat> aligned_faces;
    int face_features_length = 512;
    int face_count = 0;
    detection_face(image_input, &face_count, detection_bbox);

    if(face_count == 0){
        code = 1;  // no face
        return std::to_string(code);
    } else if(face_count > 1){
        //code = 2;  // multi face
        //return std::to_string(code);
        int first_face_size = (detection_bbox[2] - detection_bbox[0]) * (detection_bbox[3] - detection_bbox[1]);
        for(int i = 1; i < face_count; ++i){
            int face_size = (detection_bbox[i * 4 + 2] - detection_bbox[i * 4]) * (detection_bbox[i * 4 + 3] - detection_bbox[i * 4 + 1]);
            if(first_face_size < face_size){
                first_face_size = face_size;
                for(int j = 0; j < 4; ++j) detection_bbox[j] = detection_bbox[i * 4 + j];
            }
        }
        std::cout << "registration_image face_count " << face_count << ", use max size face" << std::endl;
        face_count = 1;
    }
    if(detection_bbox[2] - detection_bbox[0] < 48 || detection_bbox[3] - detection_bbox[1] < 56){
        code = 3;  // face size small
        return std::to_string(code);
    }
    float face_feature[512]; // = (float *)malloc(face_count * face_features_length *sizeof(float));
    extract_face_feature(image_input, face_count, face_feature, detection_bbox, &code, 1);
    std::string registration_image_feature_str;
    if(code == 0){
        std::string feature_str = "";
        for(int i = 0; i < 512; i++){
            feature_str += std::to_string(face_feature[i]) + ",";
        }
        registration_image_feature_str = feature_str;
    } else {
        registration_image_feature_str = std::to_string(code);
        return std::to_string(code);
    }
    if(is_recognition == 1){
        float max_score = -1.0F;
        int max_score_feature_id = -1;
        float second_score = -1.0F;
        int second_score_feature_id = -1;
        float third_score = -1.0F;
        int third_score_feature_id = -1;
        float four_score = -1.0F;
        int four_score_feature_id = -1;
        float five_score = -1.0F;
        int five_score_feature_id = -1;

        for(std::map<int, float*>::iterator it = all_feature.begin(); it != all_feature.end(); it++){
            float score = 0;
            float *register_feature = it->second;
            for(int m = 0; m < face_features_length; m++){
                score += face_feature[m] * register_feature[m];
            }
            if(score > max_score){
                five_score = four_score;
                five_score_feature_id = four_score_feature_id;
                four_score = third_score;
                four_score_feature_id = third_score_feature_id;
                third_score = second_score;
                third_score_feature_id = second_score_feature_id;
                second_score = max_score;
                second_score_feature_id = max_score_feature_id;
                max_score = score;
                max_score_feature_id = it->first;
            } else if(score > second_score){
                five_score = four_score;
                five_score_feature_id = four_score_feature_id;
                four_score = third_score;
                four_score_feature_id = third_score_feature_id;
                third_score = second_score;
                third_score_feature_id = second_score_feature_id;
                second_score = score;
                second_score_feature_id = it->first;
            } else if(score > third_score){
                five_score = four_score;
                five_score_feature_id = four_score_feature_id;
                four_score = third_score;
                four_score_feature_id = third_score_feature_id;
                third_score = score;
                third_score_feature_id = it->first;
            } else if(score > four_score){
                five_score = four_score;
                five_score_feature_id = four_score_feature_id;
                four_score = score;
                four_score_feature_id = it->first;
            } else if(score > five_score){
                five_score = score;
                five_score_feature_id = it->first;
            }
        }
        std::string return_str;
        return_str = std::to_string(max_score_feature_id) + "|" +  std::to_string(second_score_feature_id) + "|" + 
            std::to_string(third_score_feature_id) + "|" +  std::to_string(four_score_feature_id) + "|" + 
            std::to_string(four_score_feature_id);
        std::cout << "recognition_face spend: " << getMsecOfNow() - start << " ms, "
                  << max_score << " " << max_score_feature_id << " | "
                  << second_score << " " << second_score_feature_id << " | "
                  << third_score << " " << third_score_feature_id << " | "
                  << four_score << " " << four_score_feature_id << " | "
                  << five_score << " " << five_score_feature_id << std::endl;
        return return_str;
    }
    std::cout << "registration_image spend: " << getMsecOfNow() - start << " ms code " << code << ", face_count " << face_count
              << " image width " << image_input.cols << " image width " << image_input.rows
              <<  " thread_id " <<  std::this_thread::get_id() << std::endl;
    return registration_image_feature_str;
}

void sig_close(int sig){
    const char *s = NULL;
    switch (sig) {
    case SIGINT:
        s = "interrupt";
        break;
    case SIGTERM:
        s = "terminate";
        break;
    }
    recognition_stop = true;
    printf("\n\nreceived signal: %s, quit...\n", s);
}

int main(int argc, char **argv)
{
    signal(SIGINT, sig_close);  //close CTRL+C
    signal(SIGTERM, sig_close);

    rpc::server srv("127.0.0.1", 18080);
    srv.bind("registration_image", &registration_image);
    srv.bind("add_feature", add_feature);
    srv.bind("reload_feature", reload_feature);
    if(!have_init){
        init_dnn_network();
    }
    //srv.run();
    srv.async_run();
    while (!recognition_stop){
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    printf("rpc quit...\n");
    if(have_init){
        free_dnn_network();
    }
    return 0;
}
